import os
from typing import Dict, Any, Optional
from sqlalchemy import create_engine, Engine
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.ext.declarative import declarative_base
from dotenv import load_dotenv
from loguru import logger
from sqlalchemy.sql import text

# 创建基础模型
Base = declarative_base()

class DatabaseConnection:
    """数据库连接管理类，支持多种数据库类型的连接。"""
    
    def __init__(self):
        self.engine: Optional[Engine] = None
        self.session_maker = None
        # 初始化时不立即构建连接字符串，每次连接时重新获取
        self.connection_string = None
        
    def _build_connection_string(self) -> str:
        """根据环境变量构建数据库连接字符串"""
        # 强制重新加载环境变量以获得最新配置
        load_dotenv(override=True)
        
        db_type = os.getenv("DB_TYPE", "mysql")
        
        if db_type.lower() == "sqlite":
            db_path = os.getenv("DB_PATH", "database.db")
            return f"sqlite:///{db_path}"
        
        db_user = os.getenv("DB_USER", "root")
        db_password = os.getenv("DB_PASSWORD", "")
        db_host = os.getenv("DB_HOST", "localhost")
        db_port = os.getenv("DB_PORT", "3306")
        db_name = os.getenv("DB_NAME", "")
        
        logger.debug(f"数据库连接参数: 类型={db_type}, 主机={db_host}, 端口={db_port}, 用户={db_user}, 数据库={db_name}")
        
        if db_type.lower() == "mysql":
            return f"mysql+pymysql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"
        elif db_type.lower() == "postgresql":
            return f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"
        else:
            raise ValueError(f"不支持的数据库类型: {db_type}")
    
    def connect(self) -> Engine:
        """建立数据库连接并返回引擎"""
        # 每次连接时重新获取连接字符串，以反映最新的环境变量配置
        self.connection_string = self._build_connection_string()
        
        if self.engine is not None:
            # 如果引擎已存在，先关闭再重新创建
            logger.info("关闭现有数据库连接以重新连接")
            self.close()
        
        # 隐藏密码部分以增强日志安全性
        safe_connection_string = self.connection_string
        if '@' in safe_connection_string:
            parts = safe_connection_string.split('@')
            user_pass = parts[0].split('//')[-1]
            if ':' in user_pass:
                username = user_pass.split(':')[0]
                safe_connection_string = safe_connection_string.replace(user_pass, f"{username}:****")
            
        logger.info(f"正在连接到数据库: {safe_connection_string}")
        
        try:
            # 获取连接池配置
            pool_size = int(os.getenv("DB_POOL_SIZE", "5"))
            max_overflow = int(os.getenv("DB_MAX_OVERFLOW", "10"))
            pool_timeout = int(os.getenv("DB_POOL_TIMEOUT", "30"))
            
            # 根据数据库类型确定连接参数
            connect_args = {}
            db_type = os.getenv("DB_TYPE", "mysql").lower()
            
            if db_type == "mysql":
                connect_args["connect_timeout"] = 10
            
            # 尝试建立连接
            self.engine = create_engine(
                self.connection_string,
                pool_size=pool_size,
                max_overflow=max_overflow,
                pool_timeout=pool_timeout,
                connect_args=connect_args,
                echo=os.getenv("DB_ECHO", "").lower() == "true"  # 是否回显SQL
            )
            
            # 尝试连接一次以验证连接是否有效
            with self.engine.connect() as conn:
                # 获取数据库信息用于确认连接成功
                result = conn.execute(text("SELECT 1"))
                
            self.session_maker = sessionmaker(bind=self.engine)
            logger.success(f"数据库连接成功: {os.getenv('DB_HOST', 'localhost')}")
        except Exception as e:
            logger.error(f"数据库连接失败: {str(e)}")
            self.engine = None
            raise
            
        return self.engine
    
    def get_session(self) -> Session:
        """获取数据库会话"""
        if self.engine is None:
            self.connect()
        return self.session_maker()
    
    def execute_raw_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> list:
        """执行原始SQL查询"""
        with self.engine.connect() as connection:
            if params:
                result = connection.execute(query, params)
            else:
                result = connection.execute(query)
            return [dict(row) for row in result]
    
    def create_tables(self):
        """创建所有定义的表"""
        Base.metadata.create_all(self.engine)
    
    def close(self):
        """关闭数据库连接"""
        if self.engine:
            self.engine.dispose()
            logger.info("数据库连接已关闭")

# 创建全局数据库连接实例
db_connection = DatabaseConnection()

def get_db_connection() -> DatabaseConnection:
    """获取数据库连接实例"""
    return db_connection 